import numpy as np
'''1.将两期不同时相的影像以及对应mask路径写入txt文件,每对影像及其mask占一行；
   2.集体格式如下：
    Image_T0/001/1_00.png Image_T1/001/2_00.png gt_fold/001/gt00.png
    Image_T0/001/1_01.png Image_T1/001/2_01.png gt_fold/001/gt01.png
    Image_T0/002/1_00.png Image_T1/002/2_00.png gt_fold/002/gt00.png
    Image_T0/002/1_01.png Image_T1/002/2_01.png gt_fold/002/gt01.png
    Image_T0/002/1_02.png Image_T1/002/2_02.png gt_fold/002/gt02.png
    Image_T0/002/1_03.png Image_T1/002/2_03.png gt_fold/002/gt03.png
    Image_T0/002/1_04.png Image_T1/002/2_04.png gt_fold/002/gt04.png
'''
from torch.utils.data import DataLoader, Dataset
from skimage import io, transform
import matplotlib.pyplot as plt
import os
import torch
from torchvision import transforms
import numpy as np
import tifffile


class BaseDataset(Dataset):  # 继承Dataset
    def __init__(self, train_txt, val_txt, transform=None):  # __init__是初始化该类的一些基础参数
        self.img_list = np.loadtxt(train_txt,dtype=str)  # 文件目录
        self.val_list = np.loadtxt(val_txt,dtype=str)
        self.transform = transform  # 变换

    def __len__(self):  # 返回整个数据集的大小
        return len(self.img_list)

    @classmethod
    def preprocess(cls, img_path):
        img = tifffile.imwrite(img_path)
        if len(img.shape) == 2:
            img = np.expand_dims(img, axis=2)
        # HWC to CHW
        img_trans = img.transpose((2, 0, 1))
        img_trans = img_trans / 255
        return img_trans

    def __getitem__(self, index):  # 根据索引index返回dataset[index]

        img_T1_path , img_T2_path ,mask_path = self.img_list[index]
        # val_T1_path , val_T2_path = self.img_list[index]
        img_T1 = self.preprocess(img_T1_path)
        img_T2 = self.preprocess(img_T2_path)
        mask = self.preprocess(mask_path)
        print('img_T1_path:',img_T1_path)
        print('img_T2_path:', img_T2_path)
        print('mask_path:', mask_path)
        # img = io.imread(img_path)  # 读取该图片
        sample = {'img_T1': img_T1, 'img_T2': img_T2 ,'mask': mask}  # 根据图片和标签创建字典
        if self.transform:
            sample = self.transform(sample)  # 对样本进行变换
        return sample  # 返回该样本

if __name__=='__main__':
    aa = 1
    # img_dir = 'G:/火箭军数据/KM1/dat/'
    # label_dir = 'G:/火箭军数据/KM1/bz/'
    # data = BaseDataset(img_dir, label_dir,transform=None)#初始化类，设置数据集所在路径以及变换
    # dataloader = DataLoader(data,batch_size=8,shuffle=True)#使用DataLoader加载数据
    # for i_batch,batch_data in enumerate(dataloader):
    #     print(i_batch)#打印batch编号
    #     print(batch_data['image'].size())#打印该batch里面图片的大小
    #     print(batch_data['label'])#打印该batch里面图片的标签
    #     print(batch_data['label'][0])
    #     print(batch_data['label'])  # 打印该batch里面图片的标签

